#ifndef __LIBTWCLML_H__
#define __LIBTWCLML_H__

#include <linux/types.h>
#include <linux/ioctl.h>

#ifdef __cplusplus
extern "C" {
#endif

/* 设备IOCTL魔术数字 */
#define TWCLML_IOCTL_MAGIC 0x74

/* IOCTL命令定义 */
#define TWCLML_IOCTL_GET_DEVICE_COUNT _IOR(TWCLML_IOCTL_MAGIC, 0, unsigned int)
#define TWCLML_IOCTL_GET_DEVICE_INFO _IOWR(TWCLML_IOCTL_MAGIC, 1, twclml_device_info_t)
#define TWCLML_IOCTL_GET_NVLINK_INFO _IOWR(TWCLML_IOCTL_MAGIC, 2, twclml_nvlink_info_t)
#define TWCLML_IOCTL_GET_P2P_INFO _IOWR(TWCLML_IOCTL_MAGIC, 3, twclml_p2p_info_t)
#define TWCLML_IOCTL_GET_NVLINK_CAPABILITY _IOWR(TWCLML_IOCTL_MAGIC, 4, twclml_nvlink_cap_info_t)

/* 布尔类型定义 */
#ifndef bool
#define bool int
#define true 1
#define false 0
#endif

/* PCI信息结构体 */
typedef struct {
    char busId[16];                /* PCI总线ID，如"0000:01:00.0" */
    __u32 domain;                  /* PCI域 */
    __u32 bus;                     /* PCI总线号 */
    __u32 device;                  /* PCI设备号 */
    __u32 function;                /* PCI功能号 */
} nvmlPciInfo_t;

/* 启用状态枚举 */
typedef enum {
    NVML_DISABLED = 0,             /* 禁用 */
    NVML_ENABLED = 1               /* 启用 */
} nvmlEnableState_t;

/* P2P状态枚举 */
typedef enum {
    NVML_P2P_STATUS_OK = 0,                        /* 支持 */
    NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED,          /* 芯片组不支持 */
    NVML_P2P_STATUS_GPU_NOT_SUPPORTED,             /* GPU不支持 */
    NVML_P2P_STATUS_NOT_SUPPORTED,                 /* 不支持 */
    NVML_P2P_STATUS_UNKNOWN                        /* 未知 */
} nvmlGpuP2PStatus_t;

/* P2P能力索引枚举 */
typedef enum {
    NVML_P2P_CAPS_INDEX_READ = 0,      /* 读能力 */
    NVML_P2P_CAPS_INDEX_WRITE,         /* 写能力 */
    NVML_P2P_CAPS_INDEX_NVLINK,        /* NVLink P2P能力 */
    NVML_P2P_CAPS_INDEX_ATOMICS,       /* 原子操作能力 */
    NVML_P2P_CAPS_INDEX_PROP,          /* 属性 */
    NVML_P2P_CAPS_INDEX_UNKNOWN        /* 未知 */
} nvmlGpuP2PCapsIndex_t;

/* NVLink能力枚举 */
typedef enum {
    NVML_NVLINK_CAP_P2P_SUPPORTED = 0,     /* P2P支持 */
    NVML_NVLINK_CAP_SYSMEM_ACCESS = 1,     /* 系统内存访问 */
    NVML_NVLINK_CAP_P2P_ATOMICS = 2,       /* P2P原子操作 */
    NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3,    /* 系统内存原子操作 */
    NVML_NVLINK_CAP_SLI_BRIDGE = 4,        /* SLI桥接 */
    NVML_NVLINK_CAP_VALID = 5,             /* 链路有效 */
    NVML_NVLINK_CAP_COUNT                  /* 能力总数 */
} nvmlNvLinkCapability_t;

/* 错误码枚举 */
typedef enum {
    NVML_SUCCESS = 0,              /* 成功 */
    NVML_ERROR_UNINITIALIZED = 1,  /* 未初始化 */
    NVML_ERROR_INVALID_ARGUMENT = 2, /* 无效参数 */
    NVML_ERROR_NOT_SUPPORTED = 3,  /* 不支持 */
    NVML_ERROR_NOT_FOUND = 4,      /* 未找到 */
    NVML_ERROR_UNKNOWN = 999       /* 未知错误 */
} nvmlReturn_t;

/* 驱动设备信息结构体 */
typedef struct {
    unsigned int index;               /* 设备索引 */
    char bus_id[16];                  /* PCI总线ID */
    unsigned int domain;              /* PCI域 */
    unsigned int bus;                 /* PCI总线号 */
    unsigned int device;              /* PCI设备号 */
    unsigned int function;            /* PCI功能号 */
    int compute_major;                /* 计算能力主版本 */
    int compute_minor;                /* 计算能力次版本 */
    unsigned int num_nvlinks;         /* NVLink数量 */
} twclml_device_info_t;

/* 驱动NVLink信息结构体 */
typedef struct {
    unsigned int device_index;        /* 设备索引 */
    unsigned int link;                /* NVLink链路号 */
    bool is_active;                   /* 链路是否激活 */
    char remote_bus_id[16];           /* 远程设备PCI总线ID */
    unsigned int remote_domain;       /* 远程设备PCI域 */
    unsigned int remote_bus;          /* 远程设备PCI总线号 */
    unsigned int remote_device;       /* 远程设备PCI设备号 */
    unsigned int remote_function;     /* 远程设备PCI功能号 */
} twclml_nvlink_info_t;

/* 驱动P2P信息结构体 */
typedef struct {
    unsigned int device1_index;       /* 设备1索引 */
    unsigned int device2_index;       /* 设备2索引 */
    bool read_supported;              /* 读操作支持 */
    bool write_supported;             /* 写操作支持 */
    bool nvlink_supported;            /* NVLink支持 */
    bool atomics_supported;           /* 原子操作支持 */
} twclml_p2p_info_t;

/* 驱动NVLink能力信息结构体 */
typedef struct {
    unsigned int device_index;        /* 设备索引 */
    unsigned int link;                /* NVLink链路号 */
    nvmlNvLinkCapability_t capability; /* 要查询的能力 */
    unsigned int cap_result;          /* 能力查询结果 */
} twclml_nvlink_cap_info_t;

/* 设备句柄结构体 */
typedef struct twcl_nvmlDevice {
    unsigned int index;               /* 设备索引 */
    nvmlPciInfo_t pciInfo;            /* PCI信息 */
    int computeMajor;                 /* 计算能力主版本 */
    int computeMinor;                 /* 计算能力次版本 */
    unsigned int numNvLinks;          /* NVLink数量 */
} *nvmlDevice_t;

/* 动态库函数声明 */
nvmlReturn_t nvmlInit(void);
nvmlReturn_t nvmlInit_v2(void);
nvmlReturn_t nvmlShutdown(void);
nvmlReturn_t nvmlDeviceGetCount(unsigned int* count);
nvmlReturn_t nvmlDeviceGetCount_v2(unsigned int* count);
nvmlReturn_t nvmlDeviceGetHandleByIndex(unsigned int index, nvmlDevice_t* device);
nvmlReturn_t nvmlDeviceGetPciInfo_v2(nvmlDevice_t device, nvmlPciInfo_t* pci);
nvmlReturn_t nvmlDeviceGetCudaComputeCapability(nvmlDevice_t device, int* major, int* minor);
nvmlReturn_t nvmlDeviceGetNvLinkState(nvmlDevice_t device, unsigned int link, nvmlEnableState_t* isActive);
nvmlReturn_t nvmlDeviceGetNvLinkRemotePciInfo(nvmlDevice_t device, unsigned int link, nvmlPciInfo_t* pci);
nvmlReturn_t nvmlDeviceGetP2PStatus(nvmlDevice_t device1, nvmlDevice_t device2, nvmlGpuP2PCapsIndex_t p2pIndex, nvmlGpuP2PStatus_t* p2pStatus);
const char* nvmlErrorString(nvmlReturn_t r);
nvmlReturn_t nvmlDeviceGetHandleByPciBusId(const char* pciBusId, nvmlDevice_t* device);
nvmlReturn_t nvmlDeviceGetIndex(nvmlDevice_t device, unsigned int* index);
nvmlReturn_t nvmlDeviceGetNvLinkCapability(nvmlDevice_t device, unsigned int link, nvmlNvLinkCapability_t capability, unsigned int* capResult);

#ifdef __cplusplus
}
#endif

#endif /* __LIBTWCLML_H__ */
    